from crowd_sim.envs.utils.agent import Agent
from crowd_sim.envs.utils.state import JointState,JointStateStack
from collections import deque


class Robot(Agent):
    def __init__(self, config, section):
        super().__init__(config, section)

    def act(self, ob):
        if self.policy is None:
            raise AttributeError('Policy attribute has to be set!')

        state = JointState(self.get_full_state(), ob)
        action, action_index = self.policy.predict(state)
        return action, action_index


class Memory_Robot(Agent):
    def __init__(self, config, section, obs_seq_len):
        super().__init__(config, section)
        self.obs_seq_len = obs_seq_len
        self.self_state_stack =deque()

    def act(self, ob_stack):
        if self.policy is None:
            raise AttributeError('Policy attribute has to be set!')
        if len(self.self_state_stack) == self.obs_seq_len:
            self.self_state_stack.popleft()
            self.self_state_stack.append(self.get_full_state())
        else:
            while len(self.self_state_stack) < self.obs_seq_len:
                self.self_state_stack.append(self.get_full_state())
        state_stack = JointStateStack(self.self_state_stack, ob_stack)
        action, action_index = self.policy.predict(state_stack)
        return action, action_index
